package com.jd.jobkeeper.transport;

import com.google.protobuf.Message;
import com.google.protobuf.MessageLite;
import com.jd.jobkeeper.utils.InvokerHandler;
import com.jd.jobkeeper.utils.ProtobufEncoder;
import com.jd.jobkeeper.utils.ProtobufDecoder;
import io.rsocket.Payload;
import io.rsocket.RSocket;
import io.rsocket.SocketAcceptor;
import io.rsocket.core.RSocketConnector;
import io.rsocket.transport.netty.client.TcpClientTransport;
import io.rsocket.util.DefaultPayload;
import io.rsocket.util.RSocketProxy;
import reactor.core.publisher.Mono;
import reactor.netty.tcp.TcpClient;

import java.nio.charset.StandardCharsets;

public class DefaultTransport implements Transport {

    private String host;
    private int port;
    protected RSocket rSocket;

    private static final ProtobufDecoder decoder = new ProtobufDecoder();
    private static final ProtobufEncoder encoder = new ProtobufEncoder();

    private HandlerRegistry handlerRegistry;

    public void setHandlerRegistry(HandlerRegistry handlerRegistry) {
        this.handlerRegistry = handlerRegistry;
    }

    public DefaultTransport(RSocket rSocket){
        this.rSocket = rSocket;
    }

    public DefaultTransport(String host, int port) {
        this.host = host;
        this.port = port;

    }

    public void connectWithPayLoad(Payload payload) {
        this.rSocket = connectWithPayLoad(this.host, this.port, payload).block();
    }

    public Mono<RSocket> connectWithPayLoad(String host, int port, Payload payload) {
        TcpClientTransport transport = TcpClientTransport.create(
                TcpClient.create()
                        .doOnDisconnected(connection -> {
                            //disconnected
                        })
                        .host(host)
                        .port(port)
                        .wiretap(false));
        return RSocketConnector.create()
                .setupPayload(payload)
                .acceptor(createSocketAcceptor())
                .connect(transport);
    }

    @Override
    public void connect() {
        this.rSocket = connect(host, port).block();
    }

    private Mono<RSocket> connect(String host, int port) {
        TcpClientTransport transport = TcpClientTransport.create(
                TcpClient.create()
                        .doOnDisconnected(connection -> {
                            //disconnected
                        })
                        .host(host)
                        .port(Integer.valueOf(port))
                        .wiretap(false));
        return RSocketConnector.create()
                .acceptor(createSocketAcceptor())
                .connect(transport);
    }



    @Override
    public <R> Mono<R> send(String route, Message request, Class<R> type) {
        Payload payload = DefaultPayload.create(request.toByteArray(), route.getBytes(StandardCharsets.UTF_8));
        return rSocket.requestResponse(payload).map(payload1 -> decoder.decode(payload1, type)).cast(type);
    }

    public Mono<Void> fireAndForget(String route, Message request) {
        Payload payload = DefaultPayload.create(request.toByteArray(), route.getBytes(StandardCharsets.UTF_8));
        return rSocket.fireAndForget(payload);
    }

    private SocketAcceptor createSocketAcceptor() {
        return (setupPayload, sendingSocket) -> {
            RSocketProxy socketProxy = new RSocketProxy(sendingSocket) {
                @Override
                public Mono<Payload> requestResponse(final Payload payload) {
                    return handleRequestResponse(payload);
                }

                @Override
                public Mono<Void> fireAndForget(Payload payload) {
                    return handleRequestResponse(payload).cast(Void.class);
                }
            };
            return Mono.just(socketProxy);
        };
    }


    private Mono<Payload> handleRequestResponse(Payload payload) {
        try {
            String meta = payload.getMetadataUtf8();
            InvokerHandler invokerHandler = handlerRegistry.lookup(meta);
            Object arg = decoder.decode(payload, invokerHandler.getParameterType()[0]);
            return invokerHandler.invoke(arg, this)
                    .map(data -> DefaultPayload.create(encoder.encode((MessageLite) data)));
        } finally {
            payload.release();
        }
    }

    @Override
    public void dispose() {
        if (rSocket != null) {
            rSocket.dispose();
        }
    }

    @Override
    public boolean isDisposed() {
        return rSocket.isDisposed();
    }

    @Override
    public Mono<Void> onClose() {
        return this.rSocket.onClose();
    }
}
